import pandas as pd
from elasticsearch import Elasticsearch
from createES import search_embedding
import re


topics = ['United Nations Global Compact(Customer Code of Conduct/ Corporate)']
class ElasticSearchModule:

    def __init__(self):
        self.es = Elasticsearch(hosts=["http://user:pass@ip:port/"])
        self.index_name = "rfp-project"

    def search_by_text(self, text, docName, nums=20):
        body={
            'query': {
                'bool': {
                    'must': {
                        'match': {'text': text},
                    },
                    "filter":{
                        "match_phrase":{ "docName": docName}
                    }
                }
            }
        }
        text = re.sub("\(.*\)","", text)
        print(text, docName)
        res = self.es.search(index=self.index_name, body=body, size=nums)
        texts = []
        scores = []
        for ll in res["hits"]["hits"]:
            if ll["_score"]>=1.5:

                texts.append(ll["_source"]["text"])
                scores.append(ll["_score"])
        return texts, scores



    def search_by_embedding(self, text, docName, nums=10):
        body = {
            "knn": {
                "field": "text-vector",
                "query_vector": search_embedding(text),
                "k": 5,
                "num_candidates": 50,
                "filter": {
                    "match_phrase": {"docName": docName}
                }
            }
        }
        res = self.es.search(index=self.index_name, body=body, size=nums)
        texts = []
        scores = []
        for ll in res["hits"]["hits"]:
            texts.append(ll["_source"]["text"])
            scores.append(ll["_score"])

        return texts, scores

    def search_all(self,text, docName):
        key_results, key_scores = self.search_by_text(text, docName)
        emb_results, emb_scores = self.search_by_embedding(text, docName)


        all_result= {
            "keyword": [key_results[0], key_scores[0] ]if len(key_results)>0 else [],
            "embedding": [emb_results[0], emb_scores[0] ]if len(emb_results)>0 else [],
        }
        return all_result

import os


if __name__=="__main__":
    es_client_model = ElasticSearchModule()
    filep = ["../DATA/rule_process/{}".format(filename) for filename in os.listdir("../DATA/rule_process/") if
             filename.endswith("pdf")]
    datas = []

    for filename in filep:
        for topic in topics:
            filename = filename.split("/")[-1]
            try:
                res = es_client_model.search_all(topic, filename)
                if len(res["keyword"])>0:

                    datas.append([filename, topic, "key",res["keyword"][0], res["keyword"][1], ])
                if len(res["keyword"]) > 0:
                    res["embedding"][0], res["embedding"][1]
                    datas.append([filename, topic, "emb",res["embedding"][0], res["embedding"][1] ])
            except Exception as e:
                print(e)
                # continue
        # break
    df = pd.DataFrame(datas, columns=["filename", "topic","type", "text", "score"])
    df.to_excel("es_process.xlsx")


    pass





